import socket
import cv2
import pickle
import struct
import numpy as np
import time
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from geometry_msgs.msg import PoseStamped
from cv_bridge import CvBridge
from std_msgs.msg import Header
import threading
from scipy.spatial.transform import Rotation
import tf2_ros
from tf2_ros import TransformListener, Buffer

# 定义退出标志
exit_event = threading.Event()

# 定义发送图像数据及时间戳
def send_image_and_timestamp(client_socket, color_image, depth_image, cam2base_translation, cam2base_rotation, timestamp):
    """
    发送时间戳、RGB 图像和深度图像给服务器
    """
    # 数据打包
    data = {
        'timestamp': timestamp,
        'color': cv2.imencode('.jpg', color_image)[1].tobytes(),  # 压缩 RGB 图像
        'depth': depth_image.tobytes(),  # 发送原始深度图数据（np.uint16格式）
        'cam2base_trans': cam2base_translation,             # 发送相机到base_link的转换关系
        'cam2base_rot': cam2base_rotation             # 发送相机到base_link的转换关系
    }
    message = pickle.dumps(data)  # 序列化数据
    length = struct.pack('!I', len(message))  # 数据长度打包
    client_socket.sendall(length + message)  # 发送数据


def pose_msg_publish(pose_list, timestamp, publisher):
    # 创建并发布PoseStamped消息
    pose_msg = PoseStamped()
    pose_msg.header = Header()
    pose_msg.header.stamp = rclpy.time.Time(seconds=int(timestamp)).to_msg()
    pose_msg.header.frame_id = "base_link"

    pose_msg.pose.position.x = pose_list[0]
    pose_msg.pose.position.y = pose_list[1]
    pose_msg.pose.position.z = pose_list[2]
    pose_msg.pose.orientation.x = pose_list[3]
    pose_msg.pose.orientation.y = pose_list[4]
    pose_msg.pose.orientation.z = pose_list[5]
    pose_msg.pose.orientation.w = pose_list[6]

    publisher.publish(pose_msg)

# 接收服务器发送的pose数据
def receive_pose_from_server(client_socket, obj_pose_pub, button_pose_pub, machine_cup_pose_pub, handover_pose_pub):
    """
    接收服务器发送的pose数据，并发布到ROS话题
    """
    while not exit_event.is_set() and rclpy.ok():
        try:
            # 接收pose数据长度
            length_data = client_socket.recv(4)
            if not length_data:
                print("服务器连接已关闭")
                break

            # 获取数据长度
            message_length = struct.unpack('!I', length_data)[0]

            # 接收数据
            data = client_socket.recv(message_length)
            pose_data = pickle.loads(data)

            # 打印接收到的pose信息
            print(f"时间戳: {pose_data['timestamp']}")
            print(f"接收到的obj_pose数据: {pose_data['obj_pose']}")
            print(f"接收到的paper_pose数据: {pose_data['paper_pose']}")
            print(f"接收到的yoger_pose数据: {pose_data['yoger_pose']}")

            # 发布 obj_pose 话题
            pose_msg_publish(pose_data['obj_pose'], pose_data['timestamp'], obj_pose_pub)

            # 发布 paper_pose 话题
            pose_msg_publish(pose_data['paper_pose'], pose_data['timestamp'], button_pose_pub)

            # 发布 yoger_pose 话题
            pose_msg_publish(pose_data['yoger_pose'], pose_data['timestamp'], machine_cup_pose_pub)

        except Exception as e:
            print(f"接收pose数据时发生错误: {e}")
            break

# 发送图像及时间戳的线程
def image_send_thread(client_socket):
    try:
        while not exit_event.is_set() and rclpy.ok():
            # 等待ROS话题的图像消息
            color_image, depth_image, timestamp = capture_images()
            if color_image is None or depth_image is None:
                continue

            # 发送图像、转换关系和时间戳到服务器
            transform = tf_buffer.lookup_transform("base_footprint_link", "cali_camera_link", rclpy.time.Time())
            translation = np.array([transform.transform.translation.x,
                                transform.transform.translation.y,
                                transform.transform.translation.z])
            rotation = np.array([transform.transform.rotation.x,
                            transform.transform.rotation.y,
                            transform.transform.rotation.z,
                            transform.transform.rotation.w])

            send_image_and_timestamp(client_socket, color_image, depth_image, translation, rotation, timestamp)
    except Exception as e:
        print(f"发送图像时发生错误: {e}")

# 捕获图像的函数
def capture_images():
    """
    获取 RGB 和深度图像，添加检查机制，防止帧获取失败导致程序崩溃。
    """
    global color_img_msg, depth_img_msg
    try:
        # 获取当前时间戳
        timestamp = time.time()
        if color_img_msg is None or depth_img_msg is None:
            return None, None, None
        # 直接从 ROS 订阅的图像接收器中获取图像数据
        color_image = bridge.imgmsg_to_cv2(color_img_msg, desired_encoding="bgr8")
        depth_image = bridge.imgmsg_to_cv2(depth_img_msg, desired_encoding="16UC1")
        color_img_msg = None
        depth_img_msg = None
        print(timestamp)
        return color_image, depth_image, timestamp
    except Exception as e:
        print(f"捕获图像时发生错误: {e}")
        return None, None, None

# ROS 回调函数
def color_callback(msg):
    global color_img_msg
    color_img_msg = msg

def depth_callback(msg):
    global depth_img_msg
    depth_img_msg = msg

# 创建同步器并订阅图像话题
def create_sync_subscribers():
    rclpy.create_subscription(Image, "/camera/color/image_raw", color_callback, 1)
    rclpy.create_subscription(Image, "/camera/depth/image_raw", depth_callback, 1)

class ImageSenderNode(Node):

    def __init__(self):
        super().__init__('image_sender_node')
        self.obj_pose_pub = self.create_publisher(PoseStamped, '/obj_pose', 10)
        self.button_pose_pub = self.create_publisher(PoseStamped, '/paper_pose', 10)
        self.machine_cup_pose_pub = self.create_publisher(PoseStamped, '/yoger_pose', 10)

        # 启动接收位姿数据的线程
        self.receive_thread = threading.Thread(target=receive_pose_from_server, args=(client_socket, self.obj_pose_pub, self.button_pose_pub, self.machine_cup_pose_pub))
        self.receive_thread.daemon = True
        self.receive_thread.start()

        # 启动发送图像数据的线程
        self.image_thread = threading.Thread(target=image_send_thread, args=(client_socket,))
        self.image_thread.daemon = True
        self.image_thread.start()

    def destroy_node(self):
        exit_event.set()  # 通知线程退出
        super().destroy_node()

def main(args=None):
    rclpy.init(args=args)

    # 创建节点并启动
    image_sender_node = ImageSenderNode()

    try:
        rclpy.spin(image_sender_node)
    except KeyboardInterrupt:
        pass
    finally:
        image_sender_node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()
